package com.example.demo.all.controller;


import com.example.demo.all.enums.AiModel;
import com.example.demo.all.sse.SseEmitterCache;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.*;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.lang.Nullable;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author zwg
 * @date 2025/2/17 15:22
 */
@Slf4j
@Tag(name = "sse")
@RestController
@RequestMapping(path = "/sse")
public class SseController {



    @Autowired
    @Qualifier("zhiPuAiChatClient")
    ChatClient zhiPuAiChatClient;




    @GetMapping("/sendMsg")
    public String showMsg(String message){
        /*new Thread(() -> {
            try {
                for (int i = 0; i < 10; i++) {
                    // 向客户端发送数据
                    emitter.send("data: Event " + i + "\n\n");
                    log.info("send event {}",i);
                    Thread.sleep(1000); // 每秒发送一次
                }
                //emitter.complete(); // 发送完毕后，标记事件流完成
            } catch (Exception e) {
                emitter.completeWithError(e); // 如果出现异常，标记事件流出错
            }
        }).start();*/
        SseEmitter emitter = SseEmitterCache.getSseEmitter("");
        try {
            // 向客户端发送数据
            emitter.send("data: " + message + "\n\n");
            log.info("send event {}",message);
            //emitter.complete(); // 发送完毕后，标记事件流完成
        } catch (Exception e) {
            log.error("error：emitter complete", e);
            emitter.completeWithError(e); // 如果出现异常，标记事件流出错
        }
        return "send message ok";
    }



    @GetMapping("/sse")
    public SseEmitter getSse(@RequestParam(defaultValue = "") String sseId){
        return SseEmitterCache.getSseEmitter(sseId);
    }





    @Operation(summary = "chat")
    @GetMapping("/chat")
    public String chatStreamResp(
            @RequestParam AiModel aiModel,
            @RequestParam(defaultValue = "") String system,
            @RequestParam(defaultValue = "") String assistant,
            @RequestParam(defaultValue = "") String user) {

        String result = "您的消息已经收到，在思考中。。。";
        ChatClient.ChatClientRequestSpec prompt = zhiPuAiChatClient.prompt();
        SseEmitter emitter = SseEmitterCache.getSseEmitter("");
        try {
            List<Message> msgs = new ArrayList<>();
            if(StringUtils.isNotBlank(system)){
                msgs.add(new SystemMessage(system));
            }
            if(StringUtils.isNotBlank(assistant)){
                msgs.add(new AssistantMessage(assistant));
            }
            if(StringUtils.isNotBlank(user)){
                msgs.add(new UserMessage(user));
            }
            prompt.messages(msgs);

            ChatClient.StreamResponseSpec response = prompt.stream();
            log.info("resp:=====================");
            AtomicInteger count = new AtomicInteger();
            Flux<ChatResponse> responseFlux = response.chatResponse();
            responseFlux.buffer(50).subscribe((chatResponseList)->{
                try {
                    StringBuilder message = new StringBuilder();
                    for (ChatResponse chatResponse : chatResponseList) {
                        message.append(getContentFromChatResponse(chatResponse));
                    }
                    emitter.send(message.toString());
                    log.info("resp:{} {}", count.getAndIncrement(), message.toString());
                }catch (Exception e){
                    emitter.completeWithError(e);
                }
            });
        }catch (Exception e) {
            log.error("",e);
            result = e.getMessage();
            emitter.completeWithError(e);
        }

        return result;
    }


    private static String getContentFromChatResponse(@Nullable ChatResponse chatResponse) {
        return Optional.ofNullable(chatResponse)
                .map(ChatResponse::getResult)
                .map(Generation::getOutput)
                .map(AbstractMessage::getText)
                .orElse(null);
    }


}
